Using pre-trained OCP models as ASE calculators#

For those interested in using our pretrained models for other applications, we provide an [ASE](https://wiki.fysik.dtu.dk/ase/#:~:text=The%20Atomic%20Simulation%20Environment%20(ASE,under%20the%20GNU%20LGPL%20license.)-compatible Calculator to interface with ASE’s functionality.

Download pretrained checkpoint#

We have released checkpoints of all the models on the leaderboard here. These trained models can be used as an ASE calculator for various calculations.

For this tutorial we download our current best model checkpoint: GemNet-T

!wget -q -nc https://dl.fbaipublicfiles.com/opencatalystproject/models/2021_08/s2ef/gemnet_t_direct_h512_all.pt
checkpoint_path = "gemnet_t_direct_h512_all.pt"

Using the OCP Calculator#

import os

import ase.io
import numpy as np
from ase.build import add_adsorbate, fcc100, molecule
from ase.constraints import FixAtoms
from ase.optimize import BFGS
from ocpmodels.common.relaxation.ase_utils import OCPCalculator

# Construct a sample structure
adslab = fcc100("Cu", size=(3, 3, 3))
adsorbate = molecule("C3H8")
add_adsorbate(adslab, adsorbate, 3, offset=(1, 1))
tags = np.zeros(len(adslab))
tags[18:27] = 1
tags[27:] = 2
adslab.set_tags(tags)
cons = FixAtoms(indices=[atom.index for atom in adslab if (atom.tag == 0)])
adslab.set_constraint(cons)
adslab.center(vacuum=13.0, axis=2)
adslab.set_pbc(True)

# Define the calculator
calc = OCPCalculator(checkpoint=checkpoint_path)

# Set up the calculator
adslab.calc = calc

os.makedirs("data/sample_ml_relax", exist_ok=True)
opt = BFGS(adslab, trajectory="data/sample_ml_relax/toy_c3h8_relax.traj")

opt.run(fmax=0.05, steps=100)
amp: false
cmd:
  checkpoint_dir: /home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ml-catalysis-tutorials/notes/checkpoints/2022-10-30-17-23-12
  commit: cba9fb6
  identifier: ''
  logs_dir: /home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ml-catalysis-tutorials/notes/logs/tensorboard/2022-10-30-17-23-12
  print_every: 100
  results_dir: /home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ml-catalysis-tutorials/notes/results/2022-10-30-17-23-12
  seed: null
  timestamp_id: 2022-10-30-17-23-12
dataset: null
gpus: 0
logger: tensorboard
model: gemnet_t
model_attributes:
  activation: silu
  cbf:
    name: spherical_harmonics
  cutoff: 6.0
  direct_forces: true
  emb_size_atom: 512
  emb_size_bil_trip: 64
  emb_size_cbf: 16
  emb_size_edge: 512
  emb_size_rbf: 16
  emb_size_trip: 64
  envelope:
    exponent: 5
    name: polynomial
  extensive: true
  max_neighbors: 50
  num_after_skip: 2
  num_atom: 3
  num_before_skip: 1
  num_blocks: 3
  num_concat: 1
  num_radial: 128
  num_spherical: 7
  otf_graph: true
  output_init: HeOrthogonal
  rbf:
    name: gaussian
  regress_forces: true
noddp: false
optim:
  batch_size: 32
  clip_grad_norm: 10
  ema_decay: 0.999
  energy_coefficient: 1
  eval_batch_size: 32
  eval_every: 5000
  factor: 0.8
  force_coefficient: 100
  loss_energy: mae
  loss_force: l2mae
  lr_initial: 0.0005
  max_epochs: 80
  mode: min
  num_workers: 2
  optimizer: AdamW
  optimizer_params:
    amsgrad: true
  patience: 3
  scheduler: ReduceLROnPlateau
slurm: {}
task:
  dataset: trajectory_lmdb
  description: Regressing to energies and forces for DFT trajectories from OCP
  eval_on_free_atoms: true
  grad_input: atomic forces
  labels:
  - potential energy
  metric: mae
  train_on_free_atoms: true
  type: regression
trainer: forces
/home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ocp/ocpmodels/preprocessing/atoms_to_graphs.py:139: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at  /opt/conda/conda-bld/pytorch_1639180507909/work/torch/csrc/utils/tensor_new.cpp:201.)
  cell = torch.Tensor(atoms.get_cell()).view(1, 3, 3)
/home/runner/micromamba-root/envs/buildenv/lib/python3.9/site-packages/torch/autocast_mode.py:141: UserWarning: User provided device_type of 'cuda', but CUDA is not available. Disabling
  warnings.warn('User provided device_type of \'cuda\', but CUDA is not available. Disabling')
/home/runner/micromamba-root/envs/buildenv/lib/python3.9/site-packages/torch/functional.py:1069: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  /opt/conda/conda-bld/pytorch_1639180507909/work/aten/src/ATen/native/TensorShape.cpp:2157.)
  return _VF.cartesian_prod(tensors)  # type: ignore[attr-defined]
/home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ocp/ocpmodels/models/gemnet/gemnet.py:373: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').
  neighbors_new // 2,
/home/runner/work/ml_catalysis_tutorials/ml_catalysis_tutorials/ocp/ocpmodels/models/gemnet/gemnet.py:467: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').
  block_sizes = neighbors // 2
      Step     Time          Energy         fmax
BFGS:    0 17:23:40       -4.099784        1.5675
BFGS:    1 17:23:40       -4.244461        1.1370
BFGS:    2 17:23:41       -4.403120        0.7635
BFGS:    3 17:23:42       -4.503653        0.8364
BFGS:    4 17:23:43       -4.558208        0.7339
BFGS:    5 17:23:44       -4.592068        0.4095
BFGS:    6 17:23:44       -4.619359        0.7312
BFGS:    7 17:23:45       -4.671457        0.9712
BFGS:    8 17:23:46       -4.796439        0.9212
BFGS:    9 17:23:47       -4.957980        0.9762
BFGS:   10 17:23:47       -5.109450        1.0385
BFGS:   11 17:23:48       -5.295599        1.2249
BFGS:   12 17:23:49       -5.498955        1.1272
BFGS:   13 17:23:50       -5.618106        1.0669
BFGS:   14 17:23:50       -5.737125        0.9508
BFGS:   15 17:23:51       -5.901956        0.9261
BFGS:   16 17:23:52       -6.076142        1.2739
BFGS:   17 17:23:52       -6.198369        1.2028
BFGS:   18 17:23:53       -6.250305        0.6852
BFGS:   19 17:23:54       -6.254100        0.2008
BFGS:   20 17:23:55       -6.293970        0.1779
BFGS:   21 17:23:55       -6.326323        0.2295
BFGS:   22 17:23:56       -6.324457        0.1700
BFGS:   23 17:23:57       -6.321297        0.1016
BFGS:   24 17:23:57       -6.328427        0.0847
BFGS:   25 17:23:58       -6.331800        0.0587
BFGS:   26 17:23:59       -6.332060        0.0444
True
import matplotlib.pyplot as plt
from ase.visualize.plot import animate
from matplotlib import rc

# the `index` argument corresponds to what frame of the trajectory to read in, specifiying ":" reads in the full trajectory.
traj = ase.io.read("data/sample_ml_relax/toy_c3h8_relax.traj", index=":")
anim = animate(traj, radii=0.8, rotation=("-75x, 45y, 10z"))

rc("animation", html="jshtml")
anim
../_images/pretrained_models_5_1.png
import pandas as pd
import plotly.express as px

energies = [image.get_potential_energy() for image in traj]

df = pd.DataFrame(
    {
        "Relaxation Step": range(len(energies)),
        "GemNet-T Adsorption Energy [eV]": energies,
    }
)
fig = px.line(df, x="Relaxation Step", y="GemNet-T Adsorption Energy [eV]")
fig.show()